import os
import json
import pickle
from datetime import datetime
from typing import Optional, Dict, Any
from rainforeLearn.gomoku.v2.agents.dqn_agent import GomokuDQNAgent
from rainforeLearn.gomoku.v2.configs.config import GomokuDQNConfig


class CheckpointManager:
    """检查点管理器 - 负责保存和加载训练状态"""
    
    def __init__(self, config: GomokuDQNConfig):
        self.config = config
        self._ensure_checkpoint_dir_exists()
    
    def _ensure_checkpoint_dir_exists(self) -> None:
        """确保检查点目录存在"""
        os.makedirs(self.config.checkpoint_dir, exist_ok=True)
    
    def save_checkpoint(self, trainer_state: Dict[str, Any], agent: GomokuDQNAgent, 
                      checkpoint_name: Optional[str] = None) -> str:
        """保存检查点"""
        if checkpoint_name is None:
            checkpoint_name = self._generate_checkpoint_name()
        
        # 保存训练器状态和智能体模型
        checkpoint_path = self._save_trainer_state(checkpoint_name, trainer_state)
        self._save_agent_model(checkpoint_name, agent)
        self._save_checkpoint_info(checkpoint_name, trainer_state)
        
        print(f"💾 检查点已保存: {checkpoint_name}")
        return checkpoint_path
    
    def _generate_checkpoint_name(self) -> str:
        """生成检查点名称"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return f"checkpoint_{timestamp}"
    
    def _save_trainer_state(self, checkpoint_name: str, trainer_state: Dict[str, Any]) -> str:
        """保存训练器状态"""
        checkpoint_path = os.path.join(self.config.checkpoint_dir, f"{checkpoint_name}.pkl")
        with open(checkpoint_path, 'wb') as f:
            pickle.dump(trainer_state, f)
        return checkpoint_path
    
    def _save_agent_model(self, checkpoint_name: str, agent: GomokuDQNAgent) -> None:
        """保存智能体模型"""
        agent_path = os.path.join(self.config.checkpoint_dir, f"{checkpoint_name}_agent.pth")
        agent.save(agent_path)
    
    def _save_checkpoint_info(self, checkpoint_name: str, trainer_state: Dict[str, Any]) -> None:
        """保存检查点信息"""
        info_path = os.path.join(self.config.checkpoint_dir, f"{checkpoint_name}_info.json")
        
        checkpoint_info = {
            'checkpoint_name': checkpoint_name,
            'episode': trainer_state.get('current_episode', 0),
            'timestamp': datetime.now().isoformat(),
            'best_win_rate': trainer_state.get('best_win_rate', 0.0),
            'best_episode': trainer_state.get('best_episode', 0),
            'trainer_state_file': f"{checkpoint_name}.pkl",
            'agent_model_file': f"{checkpoint_name}_agent.pth",
        }
        
        with open(info_path, 'w', encoding='utf-8') as f:
            json.dump(checkpoint_info, f, indent=2, ensure_ascii=False)
    
    def load_checkpoint(self, checkpoint_filepath: str, agent: GomokuDQNAgent) -> Dict[str, Any]:
        """加载检查点"""
        checkpoint_path = self._resolve_checkpoint_path(checkpoint_filepath)
        
        if not os.path.exists(checkpoint_path):
            raise FileNotFoundError(f"检查点文件不存在: {checkpoint_path}")
        
        # 加载训练器状态和智能体模型
        trainer_state = self._load_trainer_state(checkpoint_path)
        self._load_agent_model(checkpoint_path, agent)
        
        return trainer_state
    
    def _resolve_checkpoint_path(self, checkpoint_filepath: str) -> str:
        """解析检查点路径"""
        if not checkpoint_filepath.endswith('.pkl'):
            return os.path.join(self.config.checkpoint_dir, f"{checkpoint_filepath}.pkl")
        return checkpoint_filepath
    
    def _load_trainer_state(self, checkpoint_path: str) -> Dict[str, Any]:
        """加载训练器状态"""
        with open(checkpoint_path, 'rb') as f:
            return pickle.load(f)
    
    def _load_agent_model(self, checkpoint_path: str, agent: GomokuDQNAgent) -> None:
        """加载智能体模型"""
        checkpoint_name = os.path.basename(checkpoint_path).replace('.pkl', '')
        agent_path = os.path.join(self.config.checkpoint_dir, f"{checkpoint_name}_agent.pth")
        
        if os.path.exists(agent_path):
            agent.load(agent_path)
        else:
            print(f"⚠️ 警告: 未找到智能体模型文件 {agent_path}")